import pickle
import random
from Config.Config import ACTNET200V13_PKL

with open(ACTNET200V13_PKL,'rb') as f:
    big = pickle.load(f)

small = big.copy()
big_d = big['database']
small_d = small['database']

train_num = 800
val_num   = 200

vids = list(big_d.keys())

labels_count_train = dict()
labels_count_val = dict()

for vid in vids:
    item = big_d[vid]
    if item['subset'] == 'training':
        label = item['annotations'][0]['label']
        labels_count_train[label] = labels_count_train.get(label,0)+1
        if labels_count_train[label] > 40:
            del small_d[vid]
        else:
            if random.choice([True,False,True,True]):
                del small_d[vid]
    elif item['subset'] == 'validation':
        label = item['annotations'][0]['label']
        labels_count_val[label] = labels_count_val.get(label,0)+1
        if labels_count_val[label] > 40:
            del small_d[vid]
        else:
            if random.choice([True,False,True,True]):
                del small_d[vid]
    else:
        del small_d[vid]

vids = list(small_d.keys())
a=0
b=0
for vid in vids:
    item = small_d[vid]
    if item['subset'] == 'validation':
        a+=1
    elif item['subset'] == 'training':
        b+=1

print(a,b)

with open('/mnt/md1/Experiments/TCN_Test1/PKLS/small_actNet200-V1-3.pkl','wb') as f:
    pickle.dump(small,f)

